#!/usr/bin/env python3
"""
CogVideo Environment Setup Script
Automatically detects and configures the GPU environment and installs necessary dependencies.
"""

import os
import sys
import subprocess
import platform
import logging
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class EnvironmentSetup:
    def __init__(self):
        self.platform = platform.system()
        self.python_version = sys.version_info
        self.cuda_available = False
        self.cuda_version = None
        
    def check_python_version(self):
        """Check Python version"""
        logger.info(f"Checking Python version: {sys.version}")
        if self.python_version.major == 3 and self.python_version.minor in [10, 11]:
            logger.info("✓ Python version is compatible.")
            return True
        else:
            logger.error("✗ Python 3.10 or 3.11 is required.")
            return False
    
    def check_cuda(self):
        """Check for CUDA environment"""
        logger.info("Checking for CUDA environment...")
        try:
            result = subprocess.run(
                ['nvidia-smi'], 
                capture_output=True, 
                text=True,
                shell=True
            )
            if result.returncode == 0:
                self.cuda_available = True
                logger.info("✓ NVIDIA GPU detected.")
                
                # Get CUDA version
                cuda_check = subprocess.run(
                    ['nvcc', '--version'],
                    capture_output=True,
                    text=True,
                    shell=True
                )
                if cuda_check.returncode == 0:
                    output = cuda_check.stdout
                    # Parse CUDA version
                    for line in output.split('\n'):
                        if 'release' in line.lower():
                            self.cuda_version = line.split('release')[-1].strip().split(',')[0]
                            logger.info(f"✓ CUDA version: {self.cuda_version}")
                            break
                return True
            else:
                logger.warning("✗ NVIDIA GPU not detected.")
                return False
        except FileNotFoundError:
            logger.error("✗ nvidia-smi not found. Please install NVIDIA drivers.")
            return False
        except Exception as e:
            logger.error(f"✗ CUDA check failed: {e}")
            return False
    
    def create_directories(self):
        """Create necessary directory structure"""
        logger.info("Creating directory structure...")
        directories = [
            'models',
            'logs',
            'cache',
            'uploads',
            'outputs',
            'config'
        ]
        
        for dir_name in directories:
            dir_path = Path(dir_name)
            dir_path.mkdir(exist_ok=True)
            logger.info(f"  ✓ Created directory: {dir_name}/")
    
    def install_dependencies(self):
        """Install Python dependencies"""
        logger.info("Installing Python dependencies...")
        
        # Create requirements.txt
        requirements = """# Core dependencies
# Add CogVideo dependencies here
"""
        
        req_file = Path('requirements.txt')
        req_file.write_text(requirements)
        logger.info("  ✓ Created requirements.txt")
        
        # Upgrade pip
        logger.info("  Upgrading pip...")
        subprocess.run([sys.executable, '-m', 'pip', 'install', '--upgrade', 'pip'], check=True)
        
        # Install PyTorch (select based on CUDA version)
        if self.cuda_available:
            if self.cuda_version and self.cuda_version.startswith('11'):
                torch_cmd = f"{sys.executable} -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118"
            elif self.cuda_version and self.cuda_version.startswith('12'):
                torch_cmd = f"{sys.executable} -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121"
            else:
                torch_cmd = f"{sys.executable} -m pip install torch torchvision torchaudio"
            
            logger.info(f"  Installing PyTorch (CUDA {self.cuda_version})...")
            subprocess.run(torch_cmd, shell=True, check=True)
        else:
            logger.info("  Installing PyTorch (CPU version)...")
            subprocess.run(f"{sys.executable} -m pip install torch torchvision torchaudio", shell=True, check=True)
        
        # Install other dependencies
        logger.info("  Installing other dependencies...")
        subprocess.run([sys.executable, '-m', 'pip', 'install', '-r', 'requirements.txt'], check=True)
        
        logger.info("✓ All dependencies installed.")
    
    def create_config_file(self):
        """Create default configuration file"""
        logger.info("Creating configuration file...")
        
        config_content = """# CogVideo Configuration File

server:
  host: 0.0.0.0
  port: 8080
  workers: 2
  reload: false
  max_upload_size: 10485760  # 10MB
  
model:
  name: THUDM/CogVideo
  variant: 
  device: cuda
  cache_dir: ./models
  use_safetensors: true
  enable_xformers: true  # Memory optimization
  
inference:
  # Default inference parameters
  num_frames: 25
  decode_chunk_size: 8
  fps: 7
  motion_bucket_id: 127
  noise_aug_strength: 0.02
  min_guidance_scale: 1.0
  max_guidance_scale: 3.0
  seed: -1  # -1 for random
  
  # Performance settings
  batch_size: 1
  num_inference_steps: 25
  height: 576
  width: 1024
  
logging:
  level: INFO
  format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
  log_dir: ./logs
  max_file_size: 10485760  # 10MB
  backup_count: 5
  
cache:
  enable: true
  cache_dir: ./cache
  max_cache_size: 10737418240  # 10GB
  ttl: 3600  # 1 hour
  
security:
  enable_auth: false
  api_key: ""
  max_requests_per_minute: 60
  allowed_origins:
    - "*"
  
monitoring:
  enable_metrics: true
  metrics_port: 9090
  
paths:
  upload_dir: ./uploads
  output_dir: ./outputs
  temp_dir: ./temp
"""
        
        config_file = Path('config/config.yaml')
        config_file.parent.mkdir(exist_ok=True)
        config_file.write_text(config_content)
        logger.info("  ✓ Created config/config.yaml")
    
    def verify_installation(self):
        """Verify that the installation was successful"""
        logger.info("\nVerifying installation...")
        
        try:
            # Test PyTorch
            import torch
            logger.info(f"  ✓ PyTorch version: {torch.__version__}")
            
            if torch.cuda.is_available():
                logger.info(f"  ✓ CUDA available: {torch.cuda.get_device_name(0)}")
                logger.info(f"  ✓ CUDA devices: {torch.cuda.device_count()}")
            else:
                logger.warning("  ⚠ CUDA not available, will use CPU mode.")
            
            # Test other key libraries
            import fastapi
            logger.info(f"  ✓ FastAPI version: {fastapi.__version__}")
            
            logger.info("\n✅ Environment setup successful!")
            return True
            
        except ImportError as e:
            logger.error(f"\n❌ Verification failed: {e}")
            return False
    
    def run(self):
        """Run the complete environment setup process"""
        logger.info("="*50)
        logger.info("Starting CogVideo Environment Setup")
        logger.info("="*50)
        
        steps = [
            ("Check Python version", self.check_python_version),
            ("Check CUDA environment", self.check_cuda),
            ("Create directory structure", self.create_directories),
            ("Install dependencies", self.install_dependencies),
            ("Create configuration file", self.create_config_file),
            ("Verify installation", self.verify_installation)
        ]
        
        for step_name, step_func in steps:
            logger.info(f"\nStep: {step_name}")
            logger.info("-" * 30)
            
            try:
                result = step_func()
                if result is False:
                    logger.warning(f"Step '{step_name}' completed with warnings.")
            except Exception as e:
                logger.error(f"Step '{step_name}' failed: {e}")
                logger.info("\nEnvironment setup did not complete. Please check the error messages.")
                return False
        
        logger.info("\n" + "="*50)
        logger.info("🎉 Environment setup complete!")
        logger.info("="*50)
        logger.info("\nNext steps:")
        logger.info("1. Run 'python download_models.py' to download models.")
        logger.info("2. Run 'python cogvideo_server.py' to start the server.")
        
        return True


if __name__ == "__main__":
    setup = EnvironmentSetup()
    success = setup.run()
    sys.exit(0 if success else 1)

